from datetime import datetime
import akshare as ak
import json
import sys, time
import pandas as pd
from src.mxtf_api import MxtfApi
from src.config import config

# akshare_api_doc_url = "https://akshare.akfamily.xyz/#"

class AkshareSpider:
    def __init__(self, tasks_json=None, updatedb=True):
        if tasks_json:
            self.tasks = json.loads(tasks_json)
        self.results = {}
        self.stock_list = None  # 股票列表
        self.stock_id_list = None  # 股票代码列表
        self.sw_ind_index_list = None  # 申万行业指数代码列表
        
        self.updatedb = updatedb
        
        self.mxtfapi = MxtfApi()
        self.method2apiargs = {
            "fetch_stock_hist":{
                "table_name":"stock_hist"},
            "fetch_mkt_index_hist":{
                "table_name":"mkt_index_hist"},
            "fetch_sw_ind_index_hist":{
                "table_name":"sw_ind_index_hist"},
            "fetch_sw_ind_comp":{
                "table_name":"sw_ind_comp"},
            "fetch_sw_comp_change_hist":{
                "table_name":"sw_comp_change_hist"}
            }

    @staticmethod
    def df2json(df):
        return df.to_json(orient='records')
    
    def add2db(self,table_name,data_json):
        if self.updatedb:
            results = self.mxtfapi.execute_api(table_name,'insert',data_json)
        else:
            results = ''
        return results
        
    def fetch_stock_price_adj_factors(self,horizon=5):
        
        #构建近期日期序列
        dates_in_horizon = pd.date_range(end=datetime.now(),periods=horizon,freq='d').strftime('%Y%m%d')
        #获取近期所有分红派息个股
        dividend_bonus_stocks = []
        for eachdate in dates_in_horizon:
            dividend_bonus_stocks.append(ak.news_trade_notify_dividend_baidu(date=eachdate))
        dividend_bonus_stocks = pd.concat(dividend_bonus_stocks)
            
        #获取所有近期分红派息个股的复权因子
        hist_adj_factor = []
        if not dividend_bonus_stocks.empty:
            #加工出股票id
            dividend_bonus_stocks['stock_id'] = dividend_bonus_stocks['交易所'].str.lower() + dividend_bonus_stocks['股票代码'] 
            #只筛选沪深两市股票
            dist_stock_list = dividend_bonus_stocks[(dividend_bonus_stocks['交易所']=='SH')|(dividend_bonus_stocks['交易所']=='SZ')]['stock_id'].unique()
            #获取个股复权因子
            for eachstockid in dist_stock_list:
                hist_adj_factor.append(self.fetch_single_stock_price_adj_factors(eachstockid))
        #转换数字格式适配数据库字段格式
        hist_adj_factor = pd.concat(hist_adj_factor,ignore_index=True)
        # if not hist_adj_factor.empty:
        #     hist_adj_factor['factor_value'] = hist_adj_factor['factor_value'].apply(lambda x: int(float(x) * 10**12))
        return self.df2json(hist_adj_factor)

    def fetch_single_stock_price_adj_factors(self,stock_id):
        '''获取复权因子'''
        #若输入6位数股票代码，转换为sz000001格式的的股票id
        if len(stock_id) == 6 and stock_id[:3] in config.stockcode2mkt:
            stock_id = config.stockcode2mkt[stock_id[:3]] + stock_id
        
        #获取前复权因子
        stock_hist_adj_factor_qfq = ak.stock_zh_a_daily(symbol=stock_id,adjust='qfq-factor')
        stock_hist_adj_factor_qfq.columns = ['date','factor_value']
        stock_hist_adj_factor_qfq['factor_type'] = 'qfq'
        #获取后复权因子
        stock_hist_adj_factor_hfq = ak.stock_zh_a_daily(symbol=stock_id,adjust='hfq-factor')
        stock_hist_adj_factor_hfq.columns = ['date','factor_value']
        stock_hist_adj_factor_hfq['factor_type'] = 'hfq'
        #拼接数据
        stock_hist_adj_factor = pd.concat([stock_hist_adj_factor_qfq,stock_hist_adj_factor_hfq], axis=0, ignore_index=True)
        stock_hist_adj_factor['stock_id'] = stock_id
        stock_hist_adj_factor['factor_value'] = stock_hist_adj_factor['factor_value'].apply(lambda x: int(float(x) * 10**12))
        return stock_hist_adj_factor
    def fetch_stock_ids(self,):
        if not self.stock_id_list:
            stock_zh_a_spot_em_df = ak.stock_zh_a_spot_em()
            self.stock_id_list = list(stock_zh_a_spot_em_df['代码'].unique())
        return self.stock_id_list
    def fetch_stocks(self,):
        if not self.stock_list:
            self.stock_list = ak.stock_zh_a_spot_em()
            self.stock_id_list = list(self.stock_list['代码'].unique())
        return self.df2json(self.stock_list)
    def get_stock_day_record(self,stock_id,startDate:datetime,endDate:datetime):
        start_date = startDate.strftime("%Y%m%d")
        end_date = endDate.strftime("%Y%m%d")
        data =  ak.stock_zh_a_hist(symbol=stock_id, period="daily", start_date=start_date, end_date=end_date, adjust="")
        return self.df2json(data)
    def fetch_stock_hist(self,startDate,endDate,stockList=None):
        
        if stockList == None:
            # 东财实时行情数据，获取股票列表
            stock_zh_a_spot_em_df = ak.stock_zh_a_spot_em()
            stock_id_list = list(stock_zh_a_spot_em_df['代码'].unique())
        else:
            stock_id_list = stockList.split(',')
        ful_query_data = []
        fail_list = []
        for query_stock_id in stock_id_list:
            try:
                #新浪数据
                #query_data = ak.stock_zh_a_daily(symbol=query_stock_id, start_date=query_st_date, end_date=query_ed_date, adjust="qfq")
                #东财数据
                query_data = ak.stock_zh_a_hist(symbol=query_stock_id, period="daily", start_date=startDate, end_date=endDate, adjust="")
                ful_query_data.append(query_data)
            except:
                fail_list += [query_stock_id]

            time.sleep(0.2)
        
        return self.df2json(pd.concat(ful_query_data))

    def fetch_mkt_index_hist(self,symbol,startDate:datetime,endDate:datetime):
        '''获取大盘指数历史行情数据，上证综指、深证成指和沪深300'''
        # index_id_list = ['sh000001','sh000300','sz399001']
        # all_stock_zh_index_daily_em_df = []
        startDate = startDate.strftime("%Y%m%d")
        endDate = endDate.strftime("%Y%m%d")
        stock_zh_index_daily_em_df = ak.stock_zh_index_daily_em(symbol=symbol,start_date=startDate, end_date=endDate)
        stock_zh_index_daily_em_df['symbol'] = symbol
        return self.df2json(stock_zh_index_daily_em_df)
        
    def fetch_sw_ind_comp(self):
        '''获取申万行业信息'''
        
        #申万一级行业信息
        sw_ind_first_info_df = ak.sw_index_first_info()
        sw_ind_first_info_df['sw_industry_layer'] = 'l1'
        #申万二级行业信息
        sw_ind_second_info_df = ak.sw_index_second_info()
        sw_ind_second_info_df['sw_industry_layer'] = 'l2' 
        #申万三级行业信息
        sw_ind_third_info_df = ak.sw_index_third_info()
        sw_ind_third_info_df['sw_industry_layer'] = 'l3'   
        #聚合一级二级三级行业信息
        sw_ind_all_info_df = pd.concat([sw_ind_first_info_df, sw_ind_second_info_df, sw_ind_third_info_df], axis=0, ignore_index=True)
        self.sw_ind_index_list = list(sw_ind_all_info_df['行业代码'].apply(lambda x: x[:6]))
    
        return self.df2json(sw_ind_all_info_df)
    
    def fetch_sw_ind_comp_change_hist(self):
        
        #获取历史申万行业成分
        hist_sw_ind_comp = ak.stock_industry_clf_hist_sw()
        
        return self.df2json(hist_sw_ind_comp)

    def fetch_sw_ind_index_hist(self):
        
        if self.sw_ind_index_list == None:
            _sw_ind_all_info_json = self.fetch_sw_ind_comp()
            
        #获取申万行业历史行情
        all_index_hist_sw_dfs = []
        for sw_index_id in self.sw_ind_index_list:
            index_hist_sw_df = ak.index_hist_sw(symbol=sw_index_id, period="day")
            all_index_hist_sw_dfs.append(index_hist_sw_df)
        all_index_hist_sw_df = pd.concat(all_index_hist_sw_dfs, ignore_index=True)

        return self.df2json(all_index_hist_sw_df)
    
    def fetch_sw_ind_index_hist_mini(self,sw_index_ids):
        #获取申万行业历史行情
        all_index_hist_sw_dfs = []
        for sw_index_id in sw_index_ids:
            index_hist_sw_df = ak.index_hist_sw(symbol=sw_index_id, period="day")
            all_index_hist_sw_dfs.append(index_hist_sw_df)
        all_index_hist_sw_df = pd.concat(all_index_hist_sw_dfs, ignore_index=True)

        return self.df2json(all_index_hist_sw_df)

    def fetch_data(self, task):
        
        task_id = task['taskId']
        method_name = task['methodName']

        try:
            method = getattr(self, method_name)
            if 'params' in task:
                params = task['params']
                response = method(**params)
            else:
                response = method()
            
            db_response = self.add2db(self.method2apiargs[method_name].get('table_name'),response)
            
            return [
                task_id,
                {"taskId": task_id,
                "methodName": method_name,
                "result": "success",
                "details": response,
                "db_response": db_response}
            ]
        except Exception as e:
            return [
                task_id,
                {"taskId": task_id,
                "methodName": method_name,
                "result": "failure",
                "details": str(e),
                "db_response": ""}
            ]

    def run_all(self):
        for task in self.tasks:
            task_id, results = self.fetch_data(task)
            self.results[task_id] = results

    def get_results(self):
        return self.results

def main():
    if len(sys.argv) < 2:
        print("Usage: python akshare_api_executor.py <tasks_json>")
        sys.exit(1)

    tasks_json = sys.argv[1]
    akspider = AkshareSpider(tasks_json)
    akspider.run_all()
    results = akspider.get_results()

    formatted_results = json.dumps([result for result in results.values()])
    print(formatted_results)

if __name__ == "__main__":
    main()
    
"""
传入tasks_json数组，可以传入一个或多个任务，taskId随便写，主要为了跟传出的数据关联
fetch_stock_daily-获取股票历史行情数据，起止日期必须，股票列表不是必须，没有出入股票列表，自动获取最新列表，全部更新
fetch_a_index_daily-获取大盘指数历史行情数据，上证综指、深证成指和沪深300
fetch_sw_ind-获取申万行业分类数据
fetch_hist_sw_ind_comp-获取申万行里成分股变动历史
fetch_sw_index_daily-获取申万行业指数历史行情数据
tasks_jon格式示例：
'[
    {
        "taskId": 1,
        "methodName": "fetch_stock_hist",
        "params": {
            "startDate": "20240101",
            "endDate": "20240131",
            "stockList": "000001,300339"
        }
    },
    {
        "taskId": 2,
        "methodName": "fetch_mkt_index_hist"
    },
    {
        "taskId": 3,
        "methodName": "fetch_sw_ind_comp"
    },
    {
        "taskId": 4,
        "methodName": "fetch_sw_ind_comp_change_hist"
    },
    {
        "taskId": 5,
        "methodName": "fetch_sw_ind_index_hist"
    }
]'
"""
